import torch
import torch.nn as nn
import torch.nn.functional as F
from Tuatara_Transfromer.attention import attention
import copy

class MultiHeadAttention_2(nn.Module):
    def __init__(self, num_heads, embed_dim, dropout_ratio=0.1):
        """
        Args:
            num_heads: 注意力头的数量 原版论文中为 8
            embed_dim: 输入token的维度, 即词嵌入的维度（d_model）
            dropout_ratio: Dropout概率
        """
        super().__init__()
        assert embed_dim % num_heads == 0, "embed_dim必须能被num_heads整除"

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads  # 每个头的维度

        # 定义Q、K、V的线性变换层
        # self.q_linear = nn.Linear(embed_dim, embed_dim)
        # self.k_linear = nn.Linear(embed_dim, embed_dim)
        # self.v_linear = nn.Linear(embed_dim, embed_dim)

        # copy下，linear层的参数相等
        linear_layer = nn.Linear(embed_dim, embed_dim)
        self.q_linear = copy.deepcopy(linear_layer)
        self.k_linear = copy.deepcopy(linear_layer)
        self.v_linear = copy.deepcopy(linear_layer)

        # 输出线性层
        # self.out_linear = nn.Linear(embed_dim, embed_dim)
        self.out_linear = copy.deepcopy(linear_layer)

        # Dropout
        self.dropout = nn.Dropout(dropout_ratio)

    def forward(self, query, key, value, mask=None):
        """
        Args:
            query: [batch_size, seq_len, embed_dim]
            key: [batch_size, seq_len, embed_dim]
            value: [batch_size, seq_len, embed_dim]
            mask: 可选，用于屏蔽无效位置（如padding）
        Returns:
            output: [batch_size, seq_len, embed_dim]
            attention_weights: [batch_size, num_heads, seq_len, seq_len]
        """
        batch_size = query.size(0)

        # 1. 输入层添加线性变换
        q_after_linear = self.q_linear(query)  # [batch_size, seq_len, embed_dim]
        K_after_linear = self.k_linear(key)  # [batch_size, seq_len, embed_dim]
        v_after_linear = self.v_linear(value)  # [batch_size, seq_len, embed_dim]

        # 2. 重塑为多头的形式 [batch_size, num_heads, seq_len, head_dim]
        # 进行多头拆分：每个头可以代表一部分词特征组成的句子（完整句子，完整词，只是部分维度），
        # 具体使用view方法进行拆分（维度重塑），多加了一个维度h，代表头数，
        # view中的参数-1代表自适应维度，计算机会根据这种变换自动计算这里的值
        # 原矩阵： 1-句子 2-词 3-维度（8）
        # 拆分后： 1-句子 2-拆分维度 3-词 4-拆分后的维度（假如之前是8，拆成两个就是4）
        # 然后对第二维和第三维进行转置操作，
        # 为了让代表句子长度维度和词向量维度能够相邻，这样注意力机制才能找到词义与句子位置的关系，
        q_multi = q_after_linear.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        k_multi = K_after_linear.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        v_multi = v_after_linear.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)

        # 3. 计算Scaled Dot-Product Attention
        scores = torch.matmul(q_multi, k_multi.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))

        # 4. 应用mask（如需要）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))

        # 5. 计算Attention权重
        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # 6. 加权求和
        attention_output = torch.matmul(attention_weights, v_multi)  # [batch_size, num_heads, seq_len, head_dim]

        # 自建方法（单头注意力计算， 等价于上面的 3 4 5 6）
        # 外部使用的方法有 mask 但是没有 dropout
        # attention_output, attention_weights = attention(q_multi, k_multi, v_multi, mask=mask, dropout=None)

        # 7. 合并多头并线性变换
        output = attention_output.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)
        output = self.out_linear(output)

        #
        return output, attention_weights


def test():
    torch.manual_seed(0)

    # 参数设置
    embed_dim = 8  # 输入维度
    num_heads = 2    # 头数
    seq_len = 3     # 序列长度
    batch_size = 2   # 批次大小

    # 初始化Multi-Head Attention
    mha = MultiHeadAttention_2(num_heads, embed_dim)

    # 随机生成输入数据（模拟batch_size=4, seq_len=10, embed_dim=512）
    sample_mat = torch.randn(batch_size, seq_len, embed_dim)
    # query = torch.randn(batch_size, seq_len, embed_dim)
    # key = torch.randn(batch_size, seq_len, embed_dim)
    # value = torch.randn(batch_size, seq_len, embed_dim)

    sample_mat = torch.arange(batch_size * seq_len * embed_dim, dtype=torch.float32).reshape(batch_size, seq_len, embed_dim)

    #
    query = key = value = sample_mat

    # 前向传播
    output, attn_weights = mha(query, key, value)

    print(output)

    print("Attention Output shape:", output.shape)          # [4, 10, 512]
    print("Attention weights shape:", attn_weights.shape)  # [4, 8, 10, 10]
   

if __name__ == '__main__':

    test()

    # 参数设置
    # embed_dim = 512  # 输入维度
    # num_heads = 8    # 头数
    # seq_len = 10     # 序列长度
    # batch_size = 4   # 批次大小
    #
    # # 初始化Multi-Head Attention
    # mha = MultiHeadAttention_2(num_heads, embed_dim)
    #
    # # 随机生成输入数据（模拟batch_size=4, seq_len=10, embed_dim=512）
    # query = torch.randn(batch_size, seq_len, embed_dim)
    # key = torch.randn(batch_size, seq_len, embed_dim)
    # value = torch.randn(batch_size, seq_len, embed_dim)
    #
    # # 前向传播
    # output, attn_weights = mha(query, key, value)
    #
    # print(output)
    #
    # print("Output shape:", output.shape)          # [4, 10, 512]
    # print("Attention weights shape:", attn_weights.shape)  # [4, 8, 10, 10]